
#include <windows.h>
#include <stdio.h>
#include <tchar.h>

bool regsvr(const char *component, bool install)
{
	HRESULT (*DllRegisterServer)();
	HRESULT (*DllUnregisterServer)();
	HRESULT ret;
	LPVOID msgbuf;

	HINSTANCE hHandle = LoadLibrary(component);
	if (hHandle == NULL) {
		printf("[X] Error loading %s.\n", component);
		return false;
	}

	(FARPROC&)DllRegisterServer   = GetProcAddress(hHandle,"DllRegisterServer");
	(FARPROC&)DllUnregisterServer = GetProcAddress(hHandle,"DllUnregisterServer");

	if (DllRegisterServer && DllUnregisterServer) {
		if (install) {
			ret = DllRegisterServer();
		}
		else {
			ret = DllUnregisterServer();
		}

		if (ret != S_OK) {
			FormatMessage( 
				FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM | FORMAT_MESSAGE_IGNORE_INSERTS,
				NULL,
				ret,
				MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
				(LPTSTR) &msgbuf,
				0,
				NULL 
			);
			_tprintf(_T("[X] %s"), msgbuf);
			LocalFree( msgbuf );
		}
		else {
			printf("%s %s successfully.\n", install ? "[+] Register" : "[-] Unregister", component);
		}

		FreeLibrary(hHandle);
		return (ret == S_OK);
	}
	else {
		printf("[X] No entry found in %s.\n", component);
		FreeLibrary(hHandle);
		return false;
	}
}

void help()
{
	printf("regsvr - register or unregister a DLL component.\n\n");
	printf("usage: regsvr [OPTION] COMPONENT\n");
	printf("OPTION\n");
	printf("  -i   install a component\n");
	printf("  -u   uninstall a component\n");
	printf("COMPONENT - the DLL name of component.\n\n");
}

int main(int argc, char **argv)
{
	char *component = NULL;
	bool install = true;

	for (int i = 1; i < argc; i++) {
		if (argv[i][0] == '-') {
			switch (argv[i][1]) {
			case 'i':
				install = true;
				break;

			case 'u':
				install = false;
				break;

			default:
				help();
				return -1;
			}
		}
		else {
			if (!component) {
				component = argv[i];
			}
			else {
				help();
				return -1;
			}
		}
	}

	if (!component) {
		help();
		return -1;
	}

	if (regsvr(component, install)) {
		return 0;
	}
	else {
		return 1;
	}
}
